 #include <iostream>
#include <cstdint>
#include <omp.h>
#include <libpmem.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <cassert>

#include "common.h"
#include "access.h"

using std::cout;
using std::endl;

void unit_read(char * start_addr, long size, long * rand_seed) {
	static long access_mask = (BUFSIZE - 1) ^ (ALIGN_SIZE- 1);

	asm volatile (
		"movabs $0xd800000000000000, %%rcx \n"
		"xor %%r11, %%r11 \n"
		"movq %[start_addr], %%xmm0 \n"			/* zmm0: read/write register */
        // read 1	
		RandLFSR64
		"lea (%[start_addr], %%r8), %%r9 \n"
		"xor %%r10, %%r10 \n"
"LOOP_UNIT_READ: \n"
		LD_AVX512
		"cmp %[accesssize], %%r10 \n"
		"jl LOOP_UNIT_READ \n"

		: [random]"=r"(rand_seed)
		: [start_addr]"r"(start_addr), [accesssize]"r"(size), "0"(rand_seed), [accessmask]"r"(access_mask)
		: "%rcx", "%r12", "%r11", "%r10", "%r9", "%r8"
	);
}

void unit_write(char * start_addr, long size, long * rand_seed) {
	static long access_mask = (BUFSIZE - 1) ^ (ALIGN_SIZE- 1);

	asm volatile (
		"movabs $0xd800000000000000, %%rcx \n"
		"xor %%r11, %%r11 \n"
		"movq %[start_addr], %%xmm0 \n"			/* zmm0: read/write register */
        // write 1
		RandLFSR64
		"lea (%[start_addr], %%r8), %%r9 \n"
		"xor %%r10, %%r10 \n"
"LOOP_RRW12_WR1: \n"
		WR_AVX512
		"cmp %[accesssize], %%r10 \n"
		"jl LOOP_RRW12_WR1 \n"
		BW_FENCE

		: [random]"=r"(rand_seed)
		: [start_addr]"r"(start_addr), [accesssize]"r"(size), "0"(rand_seed), [accessmask]"r"(access_mask)
		: "%rcx", "%r12", "%r11", "%r10", "%r9", "%r8"
	);
}

void print_help(char * name) {
    cout << "Usage: " << name <<" options" << endl;
    cout << "\t -p: Optane(0) or DRAM (1), default 0" << endl;
    cout << "\t -c: Cluster size, default 1" << endl;
    cout << "\t -t: Thread number, default 1" << endl;
    cout << "\t -g: Request unit, default 64" << endl;
    cout << "\t -r: Read ratio, default 0.5, value from [0, 1]" << endl;
    return ;
}

double run_test(char * buf, int size, int thread_cnt, int cluster_size, float read_ratio) {
    auto start = seconds();
    /* use Bitmap to mark the type of each operation */
    long count = BUFSIZE / size;
    long randnum = getRandom();
	long * rand_seed = &randnum;

    omp_set_num_threads(thread_cnt);
    //#pragma omp parallel for schedule(dynamic, 16)
    #pragma omp parallel for schedule(static)
    for(long i = 0; i < count; i++) {
        float tmp = long(i / cluster_size) * read_ratio;
        if((tmp - long(tmp)) >= read_ratio)
            unit_write(buf, size, rand_seed);
        else 
            unit_read(buf, size, rand_seed);
    }
    auto end = seconds();

    return end - start;
}

int main(int argc, char ** argv) {
    int opt_thread = 1;
    int opt_granu = 64;
    int opt_cluster = 1;
    float opt_read_ratio = 0.5;
    bool opt_pmem = false;

    static const char * optstr = "c:t:g:r:ph";
    char opt;
    while((opt = getopt(argc, argv, optstr)) != -1) {
        switch(opt) {
            case 'c':
                if(atoi(optarg) > 0) opt_cluster = atoi(optarg);
                break;
            case 't':
                if(atoi(optarg) > 0) opt_thread = atoi(optarg);
                break;
            case 'g':
                if(atoi(optarg) > 0) opt_granu = (atoi(optarg) / 64) * 64;;
                break;
            case 'p': 
                opt_pmem = true; 
                break;
            case 'r':
                opt_read_ratio = atof(optarg); 
                break;
            case 'h':
            case '?': print_help(argv[0]); exit(-1);;
            default:  print_help(argv[0]); exit(-1);;
        }
    }

    // print the options
    #ifdef DEBUG
        cout << "cluster size: " << opt_cluster << endl;
        cout << "read ratio: " << opt_read_ratio << endl;
        cout << "thread number: " << opt_thread << endl;
        cout << "request unit : " << opt_granu << "B" << endl;
        if(opt_pmem) {
            cout << "test environment: Optane" << endl;
        } else {
            cout << "test environment: DRAM" << endl;
        }
    #endif

    // allocating memory
    char * buf0 = NULL, *buf;
    size_t mapped_len;
    if (opt_pmem) {
        int is_pmem;
        if(file_exist("/mnt/pmem/bandwidth.pool")) {
            if(BUFSIZE > file_size("/mnt/pmem/bandwidth.pool")) {
                cout << "\nThe existing file is too small to do the test" << endl;
                exit(1);
            }
            // open the existing pool file and mmap it into virtual memory
            if ((buf0 = (char *)pmem_map_file("/mnt/pmem/bandwidth.pool", 0, 0,
                        0666, &mapped_len, &is_pmem)) == NULL) {
                perror("pmem_map_file");
                exit(1);
            }
        } else {
            // create a new pool file and mmap it into virtual memory 
            if ((buf0 = (char *)pmem_map_file("/mnt/pmem/bandwidth.pool", BUFSIZE + 1024 * 1024, PMEM_FILE_CREATE,
                        0666, &mapped_len, &is_pmem)) == NULL) {
                perror("pmem_map_file");
                exit(1);
            }
            
            // an empty file must be initialized with contents before reading
            omp_set_num_threads(16);
            #pragma omp parallel for
                for(int64_t i = 0; i < (BUFSIZE + 1024 * 1024) / 64; i++) {
                    buf0[i * 64] = 0x5a;
                }
        }

        // whether the pool file locates on PMEM or normal file system? 
        if(is_pmem == 0) {
            cout << "This file is not on pmem" << endl;
            exit(1);
        }
        buf = buf0;
    } else {
        buf0 = (char *) malloc(BUFSIZE + 1024 * 1024);
        uint64_t tmp = (uint64_t)buf0;
        uint64_t remaining = tmp & (ALIGN_SIZE - 1);
        buf = (char *) (remaining == 0 ? tmp : tmp + ALIGN_SIZE - remaining);
    }
    
    // test starts
    const int REPEAT = 2;
    double total_time = 0;
    for(int i = 0; i < REPEAT; i++) {
        drop_cache();
        total_time += run_test(buf, opt_granu, opt_thread, opt_cluster, opt_read_ratio);
    }
    cout << 64 * REPEAT / total_time << endl;

    // free memory
    if (opt_pmem)
        pmem_unmap(buf0, mapped_len);
    else
        free(buf0);
    
    return 0;
}